The function create_splash_attention_fn currently returns dot_product_attention when Splash Attention is disabled or unavailable. This is inconsistent with the new test test_create_fn_returns_none_when_disabled in tests/test_splash_attention.py, which expects None. Returning None would make the function's behavior more explicit, align with the test's expectation, and simplify the logic at call sites.

I suggest changing the implementation to return None when splash attention is not used. You'll also need to update the function's return type hint and docstring. This change will make the test test_create_fn_returns_none_when_disabled pass as written.

The initialization of self.attn can be simplified. Instead of creating a dictionary attn_kwargs and conditionally adding attention_fn, you can determine attention_fn first and then pass it directly to the nnx.MultiHeadAttention constructor. Since nnx.MultiHeadAttention handles attention_fn=None by using the default attention, the logic can be made more concise and readable.

Assuming the suggested change in splash_attention.py is made (where create_splash_attention_fn returns None if disabled), you can refactor this block to be cleaner.

        attention_fn = None
        if splash_attention_config is not None:
            attention_fn = create_splash_attention_fn(
                splash_attention_config,
                num_heads=num_heads,
                head_dim=hidden_size // num_heads,
            )

        self.attn = nnx.MultiHeadAttention(
            num_heads=num_heads,
            in_features=hidden_size,
            broadcast_dropout=False,
            decode=False,
            deterministic=False,
            dtype=dtype,
            param_dtype=param_dtype,
            rngs=rngs,
            kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), mesh_rules("map_attn_in", "map_attn_out")),
            bias_init=nnx.with_partitioning(
                nnx.initializers.zeros_init(),
                mesh_rules(
                    "map_attn_out",
                ),
            ),
            attention_fn=attention_fn,
        )

The function should return None when Splash Attention is disabled or unavailable. This allows nnx.MultiHeadAttention to fall back to its default dot_product_attention. The current implementation returns dot_product_attention directly, which is inconsistent with the new test test_create_fn_returns_none_when_disabled and can lead to less clear control flow in the model definitions.

The return type hint and docstring for create_splash_attention_fn should be updated to reflect that it can return None when splash attention is disabled. This makes the function signature more accurate and consistent with the proposed implementation change.

The logic for determining attention_fn can be simplified and made more consistent with the pattern used in src/jimm/common/vit.py. By initializing attention_fn to None and only creating the splash function if splash_attention_config is provided, the code becomes more explicit and readable. nnx.MultiHeadAttention will correctly use its default when attention_fn is None.

This change depends on create_splash_attention_fn returning None when disabled, as suggested in another comment.

        attention_fn = None
        if splash_attention_config is not None:
            attention_fn = create_splash_attention_fn(
                splash_attention_config,
                num_heads=num_heads,
                head_dim=hidden_size // num_heads,
            )

The import of dot_product_attention is performed inside this function. It's generally better practice to place all imports at the top of the file for improved readability and to avoid repeated import operations. Since flax is a core dependency of the project, you can add from flax.nnx.nn.attention import dot_product_attention at the top of the file and simplify this block.

The lambda function used here is redundant. jax.vmap can directly operate on the kernel function, which makes the code more concise and easier to read.

The logic for creating and assigning the attention_fn can be simplified. You can directly create and add the attention_fn to the attn_kwargs dictionary within a single conditional block, which removes the need for the intermediate attention_fn variable and a second if statement. This makes the code more direct and readable.

        attn_kwargs: dict = {
            "num_heads": num_heads,
            "in_features": hidden_size,
            "broadcast_dropout": False,
            "decode": False,
            "deterministic": False,
            "dtype": dtype,
            "param_dtype": param_dtype,
            "rngs": rngs,
            "kernel_init": nnx.with_partitioning(nnx.initializers.xavier_uniform(), mesh_rules("map_attn_in", "map_attn_out")),
            "bias_init": nnx.with_partitioning(
                nnx.initializers.zeros_init(),
                mesh_rules(
                    "map_attn_out",
                ),
            ),
        }
        if splash_attention_config is not None:
            attn_kwargs["attention_fn"] = create_splash_attention_fn(
                splash_attention_config,
                num_heads=num_heads,
                head_dim=hidden_size // num_heads,
            )
        self.attn = nnx.MultiHeadAttention(**attn_kwargs)

The name of this test function, test_create_fn_returns_none_when_disabled, is a bit misleading as the function doesn't return None but rather the default dot_product_attention function. Renaming it would improve clarity.